import torch
import torch.nn as nn  # 导入 PyTorch 的神经网络模块。
import torch.optim as optim  # 导入 PyTorch 的优化器模块。
from torch.utils.data import DataLoader  # 导入 PyTorch 的数据加载器。
from torchvision import datasets, transforms  # 导入 torchvision 中的数据集和数据变换模块。

# 定义超参数
batch_size = 64  # 每个批次的样本数量。
learning_rate = 0.001  # 学习率，用于优化器。
num_epochs = 10  # 训练的轮数。
model_path = 'fashionmnist_model.pth'  # 定义模型的保存路径及名称

# 数据预处理（Compose组合多个变换操作。1. ToTensor(): 将图像数据转换为 PyTorch 张量。2. 对图像数据进行归一化，均值为 0.5，标准差为 0.5。）
transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.5,), (0.5,))])  # 组合多个变换操作。

# 加载 FashionMNIST 数据集
# （root='./data': 数据集存储的路径。train=True: 是否加载训练集。download=True: 如果数据集不存在，则下载。transform=transform: 应用前面定义的数据变换。）
train_dataset = datasets.FashionMNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.FashionMNIST(root='./data', train=False, download=True, transform=transform)

# DataLoader(...): 创建数据加载器。
# train_dataset: 训练数据集。
# batch_size=batch_size: 每个批次的样本数量。
# shuffle=True: 在每个 epoch 开始时打乱数据。
# test_dataset: 测试数据集。
# shuffle=False: 测试时不打乱数据。
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)



# 定义模型
class FashionMNISTModel(nn.Module):
    def __init__(self):
        super(FashionMNISTModel, self).__init__()  # 调用父类的构造函数。
        self.fc1 = nn.Linear(28 * 28, 128)  # 定义第一个全连接层，输入维度为 28*28，输出维度为 128。
        self.fc2 = nn.Linear(128, 64)  # 定义第二个全连接层，输入维度为 128，输出维度为 64。
        self.fc3 = nn.Linear(64, 10)  # 定义第三个全连接层，输入维度为 64，输出维度为 10（对应 10 个类别）。

    def forward(self, x):
        x = x.view(-1, 28 * 28)  # 将输入张量展平为一维向量。
        x = torch.relu(self.fc1(x))  # 通过第一个全连接层并应用 ReLU 激活函数。
        x = torch.relu(self.fc2(x))  # 通过第二个全连接层并应用 ReLU 激活函数。
        x = self.fc3(x)  # 通过第三个全连接层。
        return x


model = FashionMNISTModel()  # 实例化模型。

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()  # 定义交叉熵损失函数。
optimizer = optim.Adam(model.parameters(), lr=learning_rate)  # 定义 Adam 优化器，传入模型参数和学习率。


# 训练模型
def train_model(model, train_loader, criterion, optimizer, num_epochs, model_path):
    """    定义训练模型的函数    """
    model.train()  # 设置模型为训练模式。
    for epoch in range(num_epochs):  # 迭代每个 epoch。
        running_loss = 0.0  # 初始化每个 epoch 的累计损失。
        for images, labels in train_loader:  # 迭代每个批次的数据。
            optimizer.zero_grad()  # 清除梯度。
            outputs = model(images)  # 前向传播，获取模型的输出。
            loss = criterion(outputs, labels)  # 计算损失。
            loss.backward()  # 反向传播，计算梯度。
            optimizer.step()  # 更新模型参数。
            running_loss += loss.item()  # 累计损失。
        print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {running_loss / len(train_loader):.4f}')  # 打印每个 epoch 的平均损失。
    # 保存模型
    torch.save(model.state_dict(), model_path)
    print(f'Model saved to {model_path}')

# 测试模型
def test_model(model, test_loader):
    """    定义测试模型的函数    """
    model.eval()  # 设置模型为评估模式。
    correct = 0  # 初始化正确预测的数量。
    total = 0  # 初始化总样本数量。
    with torch.no_grad():  # 关闭梯度计算，减少内存消耗。
        for images, labels in test_loader:  # 迭代每个批次的数据。
            outputs = model(images)  # 前向传播，获取模型的输出
            _, predicted = torch.max(outputs.data, 1)  # 获取每个样本的预测类别。
            total += labels.size(0)  # 累计总样本数量。
            correct += (predicted == labels).sum().item()  # 累计正确预测的数量。
    print(f'Accuracy of the model on the test images: {100 * correct / total:.2f}%')  # 打印模型的准确率。


# 训练模型
train_model(model, train_loader, criterion, optimizer, num_epochs, model_path)  # 调用训练函数。

# 测试模型 重新加载并测试
def load_model(model_path):
    model = FashionMNISTModel()
    model.load_state_dict(torch.load(model_path))
    model.eval()
    return model

# 导入模型
model = load_model(model_path)

# 测试模型
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
test_model(model, test_loader)
